import json
from datetime import datetime, timedelta
from utils.diamond131 import SWAN
import numpy as np

# 组合反射率区域: latmin, latmax, lonmin, lonmax = 19, 28, 102, 114
# apcp全区域:  18.21, 27.8, 102.2, 113.79
# 实际选取区域: 19.01, 27.8, 102.2, 113.79
latmax, latmin, lonmin, lonmax = 28, 19.01, 102, 113.99  # [20:, 20:-20]
xdim, ydim, reso = 900, 1200, 0.01

xidxs, yidxs = [], []


def get_index(file):
    with open(file, encoding='gbk') as sf:
        sdls = sf.readlines()[4:]

    for sd in sdls:
        sid, lon, lat, _, rain = sd.split()
        if latmin < float(lat) <= latmax and lonmin <= float(lon) < lonmax:
            x = round((latmax - float(lat)) / reso)
            y = round((float(lon) - lonmin) / reso)
            if x >= xdim or y >= ydim:
                continue
            xidxs.append(x)
            yidxs.append(y)
    with open("station_ids.json", "w") as jf:
        indexes = [xidxs, yidxs]
        jf.write(json.dumps(indexes))


def ts(x, y):
    assert x.shape == y.shape  # x预报，y实况
    ths = [0.1, 1, 5, 10, 15, 20]

    scores = np.zeros(len(ths))  # TS评分
    for i, th in enumerate(ths):
        # a = np.sum((x >= th)[idx] & (y >= th)[idx])  # 预报正确 [..., idx[0], idx[1]]
        a = np.sum(((x >= th) & (y >= th)))  # 预报正确
        b = np.sum(((x >= th) & (y < th)))  # 空报
        c = np.sum(((x < th) & (y >= th)))  # 漏报
        d = np.sum(((x < th) & (y < th)))  # 无降水预报正确

        all = a + b + c + d
        # assert (x.size == all.sum())

        if a + b + c > 0:
            _ts = a / (a + b + c)
            # print(_ts)
            scores[i] = _ts
        return scores


def ts_station(x, y, idx=0, threshold=0):
    assert x.shape == y.shape  # x预报，y实况
    ths = [0.1, 1, 5, 10, 15]

    # scores = []  # TS评分
    scores = np.zeros(len(ths))
    for i, th in enumerate(ths):
        # a = np.sum((x >= th)[idx] & (y >= th)[idx])  # 预报正确 [..., idx[0], idx[1]]
        a = np.sum(((x >= th) & (y >= th))[idx])  # 预报正确
        b = np.sum(((x >= th) & (y < th))[idx])  # 空报
        c = np.sum(((x < th) & (y >= th))[idx])  # 漏报
        d = np.sum(((x < th) & (y < th))[idx])  # 无降水预报正确

        # all = a + b + c + d
        # assert (x.size == all.sum())

        if a + b + c > 0:
            _ts = a / (a + b + c)
            # print(_ts)
            scores[i] = _ts
    return scores

    # return np.array(ts)


def get_time(data_time, days=0, hours=0):
    if not isinstance(data_time, datetime):
        data_time = datetime.strptime(str(data_time)[:10], '%Y%m%d%H')
    new_time = (data_time + timedelta(days=int(days), hours=int(hours))).strftime('%Y%m%d%H')
    return new_time


def get_rad(radfile):
    a2 = 118.239
    b2 = 1.5241
    min_z = -5.
    max_z = 55.

    f = SWAN(radfile)
    data = f.get_data()

    lat, lon, rad = data.variables.values()
    rad = np.minimum(max_z, rad)
    rad = np.maximum(min_z, rad)

    rad_r = np.power(10, (rad - 10 * np.log10(a2)) / (10 * b2))
    return np.nan_to_num(rad_r)/5, lat.values, lon.values


if __name__ == '__main__':
    station_grid_file = '../data/202110100012.000'

    get_index(station_grid_file)
